…rom Megatron-LM) (#427)

* CSA: add fused Compressor forward+backward CuTe-DSL kernels

Port the fused CSA/HCA Compressor gated-pooling kernels from Megatron-LM
(NVIDIA/Megatron-LM#5984) into the FE-OSS Python API,
per the maintainer request in
NVIDIA/Megatron-LM#5984 (comment).

One forward and one backward CuTe-DSL kernel fuse the THD gated-softmax
pooling region (gather -> +APE -> overlap-window transform -> fp32 softmax ->
gated weighted sum -> bf16 cast). Kernel math is unchanged from the Megatron
original (verified bitwise on forward/dKV/dScore); the interface is reshaped
to the APIBase pattern: CSACompressorForward / CSACompressorBackward classes,
csa_compressor_forward_wrapper / csa_compressor_backward_wrapper, a new
python/cudnn/csa package, docs, and fe_api tests (numerics vs fp32/upstream
eager references and an fp64 oracle, ragged packs, static-capacity padding,
run-to-run determinism, CUDA-graph capture/replay, check_support boundaries).

Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>

* CSA compressor: vectorize forward bf16 access (32-bit, CuTe autovec_copy)

Widen every bf16 load/store in the fused Compressor forward kernel from
scalar (16-bit) to paired 32-bit accesses: each thread now owns 2 adjacent
head dims and moves its window slices through 32-bit universal copies
(cute.autovec_copy over register fragments; cute.assume makes the
alignment provable). The per-lane fp32 math is unchanged and in the same
order, so the forward output stays bitwise identical to the previous
kernel (re-verified against the Megatron-LM original on the same 5 THD
shape-mix gate as the original port). Odd head_dims keep a scalar
(vec == 1) instantiation of the same kernel; the head_dim 128/512
production shapes take vec == 2.

The launch schedule moves from (rows_per_cta=4, threads=128) to one row
per CTA with 64-thread column groups, which widens the sub-wave grids
that limit the small shapes.

Measured (nsys pure kernel time, B200, ratio=4 coff=2, THD packs of
8192-token sequences, forward kernel only):

  shape        before    after    speedup
  1x8192 d128   6.0 us   4.5 us   1.36x
  3x8192 d128  12.6 us  10.0 us   1.26x
  1x8192 d512  16.0 us  12.4 us   1.29x
  3x8192 d512  42.1 us  35.0 us   1.20x

Alternatives measured and rejected on the same matrix (all bitwise-equal):
a PTX-unpack vec2 prototype (slower than the pure-DSL version:
4.7/10.5/12.9/35.7 us), and 64/128/256-bit per-thread vectors, which cut
executed instructions but lose to register pressure and occupancy
(80/147/255 regs vs 48; the 256-bit variant spills). Registers stay
spill-free at 48 (previous kernel: 40).

Tests: numerics shape matrix extended with an odd head_dim case to cover
the scalar-layout instantiation, and a check_support rejection for
head_dims beyond the launch bound.

Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>

* CSA compressor: fold dKV/dScore zero-init into the backward kernel

The backward previously required zero-initialized dKV/dScore buffers: the
kernel only wrote consumed positions, and two tensor-wide bf16 fill
kernels (torch.zeros_like) supplied the zeros everywhere else. Those
fills cost as much DRAM traffic as the kernel's own stores.

The kernel now writes every position itself: consumed positions get their
gradients as before (disjoint, atomic-free stores), and each
never-consumed slot class gets an exact zero from a unique natural owner,
keeping all stores disjoint and dKV/dScore bitwise run-to-run
deterministic:

  - first-half columns of each segment's last block's own tokens
    (coff == 2; no next block consumes them) -> that last block;
  - per-segment tail tokens (seqlen % ratio, both halves) -> the
    segment's last block;
  - all tokens of segments with zero output blocks (seqlen < ratio) ->
    CTA column bidx == 0;
  - tokens beyond cu_seqlens[-1] (static token-capacity padding of the
    gradient buffers, the CUDA-graph case) -> grid-strided across CTA
    columns.

The backward wrapper allocates grad_kv/grad_score with torch.empty_like
instead of torch.zeros_like (the fp32 dAPE fill stays: it is the atomic
accumulator). When total_comp == 0 the kernel cannot launch, so the
wrapper falls back to zeroed allocations to preserve autograd's
exact-zero semantics; the class API documents the same caveat.

Measured on the backward region (kernel + remaining fills vs kernel +
three fills before), B200, ratio=4 coff=2, THD packs of 8192-token
sequences; nsys = sum of kernel durations per iteration, wall = CUDA
events around the wrapper call:

  shape        nsys before -> after      wall before -> after
  1x8192 d128  15.5 -> 12.8 us (1.21x)   56.2 -> 47.2 us (1.19x)
  3x8192 d128  29.9 -> 22.2 us (1.35x)   61.8 -> 54.6 us (1.13x)
  1x8192 d512  34.2 -> 22.8 us (1.50x)   67.9 -> 57.3 us (1.19x)
  3x8192 d512  90.9 -> 66.0 us (1.38x)  111.6 -> 101.8 us (1.10x)

The zero-writes cost the kernel 8 registers (64 -> 72, occupancy 44% ->
41% at 3x8192 d512) and a small amount of extra store traffic, repaid
several times over by dropping the two whole-tensor fills.

dKV/dScore remain bitwise identical to the Megatron-LM original on the
5-shape port gate and to the fp32 eager reference on every test shape.

Tests: NaN-canary suite proving the kernel fully overwrites uninitialized
buffers (ragged, degenerate short segments, all-tiny packs, head_dim 512,
static-capacity padded rows) bitwise against zero-initialized runs and
the eager reference; a zeros-fallback test for packs where no segment
reaches ratio tokens. The existing CUDA-graph replay test covers the
token-capacity padding class (it fails without the grid-strided sweep).

Docs: csa.md updated (buffer contract, fill count, perf tables refreshed
for both this commit and the forward vectorization).

Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>

* Address review feedback: docs bounds, lint, warning stacklevel

- docs: list the total_comp * head_dim < 2**31 and total_comp > 0 with
  total_tokens < ratio support boundaries already enforced in check_support
- docs: run the compressor tests from test/python per repo convention
- csa/__init__.py: spell out __all__ as an explicit literal (Ruff PLE0604)
- compressor/__init__.py: sort __all__ (Ruff RUF022)
- compressor/api.py: add stacklevel=2 to the deterministic-mode RuntimeWarning

Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>

* CSA compressor: rewrite graph benchmark to fwd-only + fwd+bwd total graphs

Address review feedback on the PR #427 fairness benchmark (benchmark/csa/bench_csa_compressor.py).

- Drop the subtraction-derived backward-graph estimate. Graph variants are now two directly measured, symmetric columns for both eager and fused: a forward-only graph (eager under no_grad) and a forward+backward total graph.

- Capture the eager fwd+bwd graph with stable, pre-allocated zero .grad buffers (the supported torch CUDA-graph pattern): the captured region zeros them in place, then runs forward + autograd backward, so every replay accumulates into a zeroed buffer -- numerically identical to a single fresh backward. A per-shape graph-vs-fresh-backward cross-check is printed (all shapes bitwise-equal).

- The fused total graph captures the forward wrapper immediately followed by the backward wrapper.

- Emit both the per-call and graph tables from a single run; graph speedups are eager/fused of the displayed us, truncated to one decimal (never rounded up). The backward replay is reported only as ~total-fwd (a reference), never as a column.

Not collected by pytest; black -l160 clean.

Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>

* docs(csa): refresh graph table to fwd/total columns; soften wording

Address review feedback on the CSA compressor performance tables (docs/fe-oss-apis/csa.md).

- Replace the CUDA-graph table (which published a subtraction-derived backward column) with two directly measured columns: eager/fused forward-only graph and eager/fused forward+backward total graph. No backward-graph column; backward is noted only as ~total-fwd, never as a measured value.

- Refresh both the per-call and graph tables from a single run of the benchmark harness; graph speedups are eager/fused of the displayed us, truncated to one decimal.

- State that capture collapses each side's per-op launches into a single replay rather than removes launch/host overhead; say less launch-bound rather than compute-bound.

- Note how the re-run reproduces the previously published per-call wall clock, and document the eager-backward capture basis (stable zero .grad buffers).

Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>

* bench(csa): guard speedup division, dedupe leaf construction, ruff nits

Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>

* CSA compressor: widen the validated envelope to coff in {1, 2}

The kernels are already generic over coff — for coff=1 the window is the
block's own ratio tokens (no overlap transform, always valid); only the
check_support gate restricted the APIs to the production coff=2 form.
Widen the gate to coff in {1, 2} (ratio stays gated at 4), update the
error message and the envelope prose in api.py / docs / README, and
parametrize the tests over both coff forms: numerics vs the
fp32/upstream/fp64 eager references (the eager reference already
implements coff=1 faithfully — it skips the overlap transform), static-
capacity padding, replay determinism, empty outputs, NaN-canary
zero-writes, deterministic-mode error paths, CUDA-graph capture,
class-vs-wrapper equivalence, and check_support accept/reject
boundaries (new coff=0 / coff=3 rejects; a new empty-segment shape
covers both coff forms). No kernel changes.

Addresses hxbai's coff=1 question on the PR.

Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>

* CSA: add missing docstrings across the compressor Python files

Bring the PR's Python docstring coverage over CodeRabbit's 80% pre-merge
gate: module docstrings for the csa package/compressor __init__ files
and one-line docstrings for the remaining undocumented helpers in
api.py (cache/ctor/compile plumbing), compressor_sm100.py (launcher
ctors), the test helpers, and the benchmark helpers. interrogate 1.7.0
over the PR's six CSA Python files: 69.4% -> 100%. Docstrings only — no
logic, kernel, or launch-path changes.

Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>

* CSA compressor: add dedicated ratio=128 kernels (forward+backward) with ratio dispatch

The generic CSA compressor kernels keep the whole coff * ratio pooling
window in per-thread registers; past ratio ~32 that hits the
255-register wall and spills kilobytes of local memory per thread, so
ratio=128 needs dedicated kernels rather than a lifted gate. This adds
them behind the same public API: check_support now admits
ratio in {4, 128} x coff in {1, 2} (ratio=128 additionally gated to
head_dim in {128, 512}, the hardware-validated set) and the wrappers
and class APIs route to the matching kernel family by ratio
transparently.

Design: the forward streams the window with a chunked online softmax
(one CTA per output row, per-column (max, denom, acc) triples merged in
one fixed smem order); its launch schedule is bucketed by output-row
count (small/default/large, all precompiled for CUDA-graph safety) and
per bucket selects two-phase accumulation and an ex2.approx fast exp
where measured faster. The backward stages each row's window through
shared memory chunk-parallel, fuses the per-chunk den/S partial sums
into the same pass, merges partials in a fixed chunk order, and stores
gradients with a hoisted 1/den multiply — keeping kernel-side
zero-writes to never-consumed dKV/dScore slots and fp32-atomic dAPE
exactly as at ratio=4. ptxas (sm_100a): forward 32-51 registers,
backward 48-128, 0 spill / 0 stack across all 16 shipped
(config, schedule) kernels (reproducer:
benchmark/csa/reg_probe_csa_compressor_r128.py).

Numerics contract, per ratio family (docs/fe-oss-apis/csa.md): ratio=4
keeps its bitwise contract UNCHANGED — dKV/dScore bit-identical to the
fp32-intermediate eager reference, at coff 1 and 2. ratio=128 is
deterministic + faithful to that same fp32-intermediate eager
reference: bitwise run-to-run determinism of out/dKV/dScore
(NaN-prefill replay tested); the same values within final-bf16 rounding
at the gate tolerances (differing elements <= max(1, 0.1%), max_abs <=
1.6e-2, calibrated on the gate's documented input distribution — bf16
deviations scale 2^k with 2^k inputs while the fp32 intermediates stay
finite); the eager reference's non-finite propagation (finite inputs
that overflow its fp32 intermediates poison both sides alike, with the
stated caveat that the fused order saturates earlier near fp32 max —
un-normalized chunk partials); and, on inputs whose fp32 intermediates
stay finite in both evaluation orders, fp64-oracle parity (at least as
close to an fp64 oracle as the eager reference). The reduction reorders
and fast-exp buckets were each adopted on a measured win and are
covered by that contract.

Performance (nsys isolated GPU kernel time, 1x B200, vs the
fp32-intermediate eager reference region on identical inputs): forward
13.2-44.4x, backward 7.4-21.9x across coff {1, 2} x head_dim {128, 512}
x 8k-131k tokens; e.g. coff=2, d=128, 65k tokens: fwd 726.3 -> 16.4 us
(44.4x), bwd 960.0 -> 61.2 us (15.7x). Full tables with methodology in
docs/fe-oss-apis/csa.md.

Validation (B200, CC 10.0): the committed 21-case contract gate
(benchmark/csa/gate_csa_compressor_r128.py) passes 21/21 with shipped
schedule coverage asserted (forward 10/10, backward 6/6 buckets);
pytest fe_api/csa/test_CSA_compressor.py: 88 passed, 1 skipped at the
default L0 level and 98 passed, 1 skipped with -m "L0 or L1" (both
ratio families; the ratio=4 rows, including the coff=1 envelope, stay
on their bitwise assertions). New tests cover the ratio=128 dispatch
envelope (schedule selection at every nb_total bucket boundary; every
shipped (config, schedule) kernel executed against the full contract at
L1), exact-zero assertions on every never-consumed dKV/dScore slot
class in the NaN-canary, CUDA-graph capture/replay for the new family,
and grad_ape zeroing-ownership regressions.

Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>

---------

Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>